#!/usr/bin/env python3
"""
水文预报分析服务器

提供：
1. WebSocket 接口 - 实时水文预报分析
2. HTTP REST API - 参数修改和降雨假拟工具（新增功能）
"""

import json
import os
from pathlib import Path

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from mcp.forecast_analyzer import ForecastAnalyzer
from mcp.forecast_client import ForecastAPIClient
from routers.entity_router import router as entity_router, init_matcher
from routers.tools_router import router as tools_router

# 加载环境变量（override=True 确保容器内的 .env 优先）
env_path = Path(__file__).parent.parent / ".env"
load_dotenv(env_path, override=True)


# ==================== 配置 ====================

class Config:
    """服务配置"""
    # vLLM
    VLLM_BASE_URL = os.getenv("VLLM_BASE_URL", "http://localhost:8000/v1")
    VLLM_MODEL_NAME = os.getenv("VLLM_MODEL_NAME", "Qwen/Qwen2.5-3B-Instruct")
    VLLM_TEMPERATURE = float(os.getenv("VLLM_TEMPERATURE", "0.7"))
    VLLM_MAX_TOKENS = int(os.getenv("VLLM_MAX_TOKENS", "2048"))
    VLLM_API_KEY = os.getenv("VLLM_API_KEY", "EMPTY")

    # 服务配置（统一端口）
    HOST = os.getenv("WEBSOCKET_HOST", "0.0.0.0")
    PORT = int(os.getenv("WEBSOCKET_PORT", "8765"))

    # SSL 配置（可选）
    SSL_ENABLED = os.getenv("SSL_ENABLED", "false").lower() == "true"
    SSL_CERTFILE = os.getenv("SSL_CERTFILE", "/etc/letsencrypt/live/ws.waterism.tech/fullchain.pem")
    SSL_KEYFILE = os.getenv("SSL_KEYFILE", "/etc/letsencrypt/live/ws.waterism.tech/privkey.pem")

    # 默认 Session
    DEFAULT_SESSION_ID = os.getenv("DEFAULT_SESSION_ID", "daa0d8a6-dfb3-4887-b798-366f4682c453")


# ==================== 工具函数 ====================

async def get_session_info(session_id: str) -> str:
    """获取预报会话信息"""
    client = ForecastAPIClient()
    data = await client.fetch_forecast_data(session_id)
    analyzer = ForecastAnalyzer(data)

    session_info = analyzer.get_session_info()
    model_params = analyzer.get_model_parameters()

    result = f"""
预报会话信息：
- 会话ID: {session_info['session_id']}
- 站点: {session_info['station_name']} ({session_info['station_code']})
- 模型类型: {session_info['model_type']}
- 模型链: {session_info['model_chain']}
- 假拟预报数量: {session_info['num_hypothetical']}

模型参数：
- 模型名称: {model_params.get('model_name')}
- 初始状态: SA0={model_params.get('initial_state', {}).get('SA0')}, UA0={model_params.get('initial_state', {}).get('UA0')}, YA0={model_params.get('initial_state', {}).get('YA0')}
"""
    return result.strip()


async def analyze_metrics(forecast_type: str, session_id: str) -> str:
    """分析洪水指标"""
    client = ForecastAPIClient()
    data = await client.fetch_forecast_data(session_id)
    analyzer = ForecastAnalyzer(data)

    results = []

    if forecast_type in ["real", "all"]:
        real_analysis = analyzer.analyze_real_forecast()
        metrics = real_analysis["metrics"]

        if "error" not in metrics:
            results.append(f"""
真实预报指标：
- 洪量: {metrics['flood_volume_m3']:,.2f} m³
- 洪峰流量: {metrics['peak_discharge_m3s']:.2f} m³/s
- 峰现时间: {metrics['peak_time']}
- 平均流量: {metrics['avg_discharge_m3s']:.2f} m³/s
- 数据点数: {metrics['num_points']}
- 预报时段: {metrics['start_time']} 至 {metrics['end_time']}
""")

    if forecast_type in ["hypothetical", "all"]:
        hyp_analyses = analyzer.analyze_all_hypothetical()

        for i, hyp in enumerate(hyp_analyses):
            metrics = hyp["metrics"]
            if "error" not in metrics:
                results.append(f"""
假拟预报 #{i} ({hyp['name']}):
- 洪量: {metrics['flood_volume_m3']:,.2f} m³
- 洪峰流量: {metrics['peak_discharge_m3s']:.2f} m³/s
- 峰现时间: {metrics['peak_time']}
- 平均流量: {metrics['avg_discharge_m3s']:.2f} m³/s
""")

    return "\n".join(results).strip() if results else "未找到有效的预报数据"


async def compare_forecasts(session_id: str) -> str:
    """对比预报差异"""
    client = ForecastAPIClient()
    data = await client.fetch_forecast_data(session_id)
    analyzer = ForecastAnalyzer(data)

    comparison = analyzer.compare_forecasts()

    if "error" in comparison:
        return f"错误: {comparison['error']}"

    real_metrics = comparison["real_forecast"]
    results = [f"""
真实预报基准值：
- 洪量: {real_metrics['flood_volume_m3']:,.2f} m³
- 洪峰流量: {real_metrics['peak_discharge_m3s']:.2f} m³/s
- 峰现时间: {real_metrics['peak_time']}

假拟预报对比（共 {comparison['num_hypothetical']} 个）：
"""]

    for comp in comparison["comparisons"]:
        if "error" in comp:
            results.append(f"\n{comp['name']}: {comp['error']}")
            continue

        diff = comp["differences"]
        results.append(f"""
{comp['name']}:
  洪量差异: {diff['flood_volume']['difference_m3']:+,.2f} m³ ({diff['flood_volume']['difference_percent']:+.2f}%)
  洪峰差异: {diff['peak_discharge']['difference_m3s']:+.2f} m³/s ({diff['peak_discharge']['difference_percent']:+.2f}%)
  峰现时间: {diff['peak_time']['hypothetical']} {'(相同)' if diff['peak_time']['same'] else f'(真实: {diff["peak_time"]["real"]})'}
""")

    # 统计摘要
    summary = comparison["summary"]
    if "note" not in summary:
        results.append(f"""
统计摘要：
  洪量差异范围: {summary['volume_diff_range']['min_percent']:.2f}% ~ {summary['volume_diff_range']['max_percent']:.2f}% (平均: {summary['volume_diff_range']['avg_percent']:.2f}%)
  洪峰差异范围: {summary['peak_diff_range']['min_percent']:.2f}% ~ {summary['peak_diff_range']['max_percent']:.2f}% (平均: {summary['peak_diff_range']['avg_percent']:.2f}%)
""")

    return "".join(results).strip()


# ==================== LLM 处理 ====================

async def ask_with_llm(question: str, session_id: str) -> str:
    """使用 LLM 理解问题并调用相应工具"""

    # 初始化 LLM
    llm = ChatOpenAI(
        base_url=Config.VLLM_BASE_URL,
        model_name=Config.VLLM_MODEL_NAME,
        api_key=Config.VLLM_API_KEY,
        temperature=0.3,
        max_tokens=500,
    )

    # 分析用户问题，决定调用哪个工具
    classifier_prompt = ChatPromptTemplate.from_messages([
        ("system", """你是一个水文预报分析助手。根据用户问题，判断应该调用哪个工具：

工具1: get_session_info - 查询会话基本信息、站点、模型信息
工具2: analyze_metrics - 分析洪量、洪峰、峰现时间等指标
工具3: compare_forecasts - 对比真实预报和假拟预报的差异

请只回复工具名称，不要有其他内容。如果不确定，回复 analyze_metrics。"""),
        ("human", "{question}")
    ])

    # 让 LLM 选择工具
    chain = classifier_prompt | llm
    response = await chain.ainvoke({"question": question})
    tool_name = response.content.strip().lower()

    # 调用相应工具
    if "session" in tool_name or "info" in tool_name:
        tool_result = await get_session_info(session_id)
    elif "compare" in tool_name or "对比" in tool_name or "差异" in tool_name:
        tool_result = await compare_forecasts(session_id)
    else:
        tool_result = await analyze_metrics("all", session_id)

    # 让 LLM 用自然语言总结
    answer_prompt = ChatPromptTemplate.from_messages([
        ("system", """你是水文预报分析专家。根据工具返回的数据，用清晰专业的语言回答用户问题。

术语说明：
- SA0: 流域蓄水初值
- UA0: 地下水初值  
- YA0: 河道蓄水初值
- 洪量: 洪水总量（m³）
- 洪峰: 最大流量（m³/s）
- 峰现时间: 洪峰出现时间

请直接回答问题，不要重复工具输出。"""),
        ("human", "用户问题: {question}\n\n工具返回数据:\n{tool_result}\n\n请回答用户问题:")
    ])

    answer_chain = answer_prompt | llm
    final_response = await answer_chain.ainvoke({
        "question": question,
        "tool_result": tool_result
    })

    return final_response.content


# ==================== FastAPI 应用 ====================

# 创建 FastAPI 应用
app = FastAPI(
    title="水文预报分析服务",
    description="提供 WebSocket 实时分析和 HTTP REST API 工具服务",
    version="2.0.0"
)

# 配置 CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 引入工具路由（参数修改 + 降雨假拟）
app.include_router(tools_router)

# 引入实体匹配路由（Trie 树匹配）
app.include_router(entity_router)


# ==================== WebSocket 处理 ====================

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket 端点 - 处理实时水文预报分析"""
    await websocket.accept()
    client_id = id(websocket)
    print(f"🔌 WebSocket 客户端连接: {client_id}")

    try:
        # 发送欢迎消息
        await websocket.send_json({
            "type": "welcome",
            "message": "🌊 水文预报分析服务已连接",
            "config": {
                "model": Config.VLLM_MODEL_NAME,
                "default_session": Config.DEFAULT_SESSION_ID
            }
        })

        # 处理消息
        while True:
            try:
                # 接收消息
                data = await websocket.receive_json()

                # 提取参数
                session_id = data.get("session_id", Config.DEFAULT_SESSION_ID)
                question = data.get("question", "对比真实预报和假拟预报在洪量、洪峰、峰现时间上的差异")

                # 发送处理中状态
                await websocket.send_json({
                    "type": "processing",
                    "message": "正在分析..."
                })

                # 调用 LLM 处理
                answer = await ask_with_llm(question, session_id)

                # 返回结果
                await websocket.send_json({
                    "type": "answer",
                    "question": question,
                    "answer": answer,
                    "session_id": session_id
                })

            except json.JSONDecodeError:
                await websocket.send_json({
                    "type": "error",
                    "error": "无效的 JSON 格式"
                })
            except Exception as e:
                await websocket.send_json({
                    "type": "error",
                    "error": str(e)
                })
                print(f"❌ 处理错误: {e}")
                import traceback
                traceback.print_exc()

    except WebSocketDisconnect:
        print(f"🔌 WebSocket 客户端断开: {client_id}")
    except Exception as e:
        print(f"❌ WebSocket 连接错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        print(f"👋 WebSocket 客户端离开: {client_id}")


# ==================== HTTP 路由 ====================

@app.get("/")
async def root():
    """根路径 - 服务信息"""
    return {
        "service": "水文预报分析服务",
        "version": "2.0.0",
        "endpoints": {
            "websocket": "/ws",
            "tools_api": "/api/execute",
            "health": "/api/health",
            "docs": "/docs"
        }
    }


@app.get("/ws/info", tags=["WebSocket"])
async def websocket_info():
    """
    WebSocket 连接信息
    
    WebSocket 端点不会出现在 Swagger 文档中（OpenAPI 不支持 WebSocket 协议）
    
    ## 连接地址
    - **端点**: `ws://host:8765/ws`
    - **协议**: WebSocket
    
    ## 连接示例
    
    ### JavaScript
    ```javascript
    const ws = new WebSocket('ws://localhost:8765/ws');
    
    ws.onopen = () => {
      console.log('Connected');
      ws.send(JSON.stringify({
        session_id: "daa0d8a6-dfb3-4887-b798-366f4682c453",
        question: "对比真实预报和假拟预报的差异"
      }));
    };
    
    ws.onmessage = (event) => {
      const data = JSON.parse(event.data);
      console.log('Received:', data);
    };
    ```
    
    ### Python
    ```python
    import asyncio
    import websockets
    import json
    
    async def test():
        uri = "ws://localhost:8765/ws"
        async with websockets.connect(uri) as websocket:
            # 发送问题
            await websocket.send(json.dumps({
                "session_id": "daa0d8a6-dfb3-4887-b798-366f4682c453",
                "question": "对比真实预报和假拟预报的差异"
            }))
            
            # 接收响应
            response = await websocket.recv()
            print(json.loads(response))
    
    asyncio.run(test())
    ```
    
    ## 消息格式
    
    ### 请求格式
    ```json
    {
      "session_id": "会话ID",
      "question": "你的问题"
    }
    ```
    
    ### 响应格式
    
    #### 欢迎消息
    ```json
    {
      "type": "welcome",
      "message": "🌊 水文预报分析服务已连接",
      "config": {
        "model": "Qwen/Qwen2.5-3B-Instruct",
        "default_session": "daa0d8a6-dfb3-4887-b798-366f4682c453"
      }
    }
    ```
    
    #### 处理中
    ```json
    {
      "type": "processing",
      "message": "正在分析..."
    }
    ```
    
    #### 回答
    ```json
    {
      "type": "answer",
      "question": "用户问题",
      "answer": "AI回答内容",
      "session_id": "会话ID"
    }
    ```
    
    #### 错误
    ```json
    {
      "type": "error",
      "error": "错误信息"
    }
    ```
    
    ## 测试工具
    
    - **在线测试**: 使用浏览器 WebSocket 测试工具
    - **命令行**: `python examples/websocket_client_example.py`
    - **curl（不支持）**: curl 不支持 WebSocket 协议
    
    ## 注意事项
    
    1. WebSocket 连接是持久连接，不同于 HTTP 的请求-响应模式
    2. 连接建立后会收到欢迎消息
    3. 可以在同一连接上发送多个问题
    4. 每个问题会先收到 "processing" 消息，然后收到 "answer"
    5. 使用完毕后记得关闭连接
    """
    return {
        "endpoint": "/ws",
        "protocol": "WebSocket",
        "description": "实时水文预报分析服务",
        "url": "ws://localhost:8765/ws",
        "message_types": ["welcome", "processing", "answer", "error"],
        "example_client": "examples/websocket_client_example.py",
        "note": "WebSocket 端点不会出现在 Swagger UI 中（OpenAPI 规范不支持 WebSocket）"
    }


# ==================== 服务器启动 ====================

def main():
    """启动服务器"""
    print("=" * 80)
    print("🌊 水文预报分析服务")
    print("=" * 80)

    # 初始化实体匹配器（加载 CSV 数据）
    init_matcher()
    print(f"\n配置信息:")
    print(f"  vLLM 地址: {Config.VLLM_BASE_URL}")
    print(f"  vLLM 模型: {Config.VLLM_MODEL_NAME}")

    # 判断协议
    protocol = "https" if Config.SSL_ENABLED else "http"
    ws_protocol = "wss" if Config.SSL_ENABLED else "ws"

    if Config.SSL_ENABLED:
        print(f"  SSL 模式: 已启用 🔒")
        print(f"  SSL 证书: {Config.SSL_CERTFILE}")
        print(f"  服务地址: {protocol}://{Config.HOST}:{Config.PORT}")
    else:
        print(f"  SSL 模式: 未启用")
        print(f"  服务地址: {protocol}://{Config.HOST}:{Config.PORT}")

    print(f"  默认会话ID: {Config.DEFAULT_SESSION_ID}")
    print("\n" + "=" * 80)
    print(f"🚀 服务器启动在: {protocol}://{Config.HOST}:{Config.PORT}")
    print("=" * 80)
    print(f"\n📚 可用接口:")
    print(f"  - WebSocket: {ws_protocol}://{Config.HOST}:{Config.PORT}/ws")
    print(f"  - 工具 API: {protocol}://{Config.HOST}:{Config.PORT}/api/execute")
    print(f"  - 健康检查: {protocol}://{Config.HOST}:{Config.PORT}/api/health")
    print(f"  - API 文档: {protocol}://{Config.HOST}:{Config.PORT}/docs")
    print(f"  - WebSocket 文档: {protocol}://{Config.HOST}:{Config.PORT}/ws/info")
    print("=" * 80 + "\n")

    # 启动 FastAPI 服务器
    if Config.SSL_ENABLED:
        # SSL 模式
        uvicorn.run(
            app,
            host=Config.HOST,
            port=Config.PORT,
            ssl_certfile=Config.SSL_CERTFILE,
            ssl_keyfile=Config.SSL_KEYFILE,
            log_level="info"
        )
    else:
        # HTTP 模式
        uvicorn.run(
            app,
            host=Config.HOST,
            port=Config.PORT,
            log_level="info"
        )


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n👋 服务器已停止")
